import os
from PIL import Image, ImageDraw


def draw_and_crop(image_path, rectangles, output_folder):
    """
    Draws rectangles on an image and crops the regions defined by the rectangles.

    :param image_path: Path to the input image.
    :param rectangles: List of tuples, each containing the coordinates for a rectangle
                       (left, top, right, bottom) and the color to draw the rectangle.
    :param output_folder: Folder to save the output images.
    """
    # Ensure the output directory exists
    os.makedirs(output_folder, exist_ok=True)

    # Load the image
    image = Image.open(image_path)
    draw = ImageDraw.Draw(image)

    # Process each rectangle
    for idx, (coords, color) in enumerate(rectangles, start=1):
        # Draw the rectangle on the image
        draw.rectangle(coords, outline=color, width=2)

        # Crop the rectangle and save the cropped image
        cropped_image = image.crop(coords)
        base_filename = os.path.splitext(os.path.basename(image_path))[0]
        cropped_filename = f"{base_filename}_crop{idx}.png"
        cropped_image.save(os.path.join(output_folder, cropped_filename))

    # Save the image with drawn rectangles
    base_filename = os.path.splitext(os.path.basename(image_path))[0]
    output_image_path = os.path.join(output_folder, f"{base_filename}_with_rectangles.png")
    image.save(output_image_path)

    return output_image_path


# Example usage:
image_paths = [
    'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\wrapgt.png',
    'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\bUnwarp_wrap.png',
    'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\elastic_wrap.png',
    'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\emreg_wrap.png',
    'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\ours_wrap.png',
    'C:\\Users\\crxc\\Pictures\\wandb\\论文2\\ssemnet_wrap.png',
    # 'image2.png',
    # Add your other image paths here
]
rectangles_info = [
    ((200, 0, 300, 50), 'red'),
    ((180, 220, 290, 320), 'blue'),
    # Add other rectangles here
]
output_folder = 'C:\\Users\\crxc\\Pictures\\wandb\\论文2'

# Call the function for each image
for image_path in image_paths:
    draw_and_crop(image_path, rectangles_info, output_folder)
